import os
import time
import torch
import argparse # 用于命令行参数解析
import pandas as pd
import seaborn as sns # 这是一个数据可视化的库 绘制散点图
import matplotlib.pyplot as plt
from torchvision import transforms # 用于数据预处理
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
from collections import defaultdict # 用于记录日志

from models import VAE


def main(args):
    torch.manual_seed(args.seed)  # 设置随机种子
    if torch.cuda.is_available():  # 如果有GPU
        torch.cuda.manual_seed(args.seed)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    ts = time.time()  # 记录开始时间

    dataset = MNIST(
        root='data', train=True, transform=transforms.ToTensor(),   # 下载MNIST数据集
        download=True)
    data_loader = DataLoader(
        dataset=dataset, batch_size=args.batch_size, shuffle=True)

    def loss_fn(recon_x, x, mean, log_var):  # 定义损失函数
        BCE = torch.nn.functional.binary_cross_entropy(  # 计算BCE损失  二进制交叉熵
            recon_x.view(-1, 28 * 28), x.view(-1, 28 * 28), reduction='sum')   #将张量展平为28*28的大小
        KLD = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp())   # 计算KLD损失  正则化项

        return (BCE + KLD) / x.size(0)  # 计算总损失

    vae = VAE(
        encoder_layer_sizes=args.encoder_layer_sizes,   # 定义VAE模型
        latent_size=args.latent_size,   # 定义隐变量维度
        decoder_layer_sizes=args.decoder_layer_sizes,  # 定义解码器层数
        conditional=args.conditional,  # 是否使用条件VAE
        num_labels=10 if args.conditional else 0).to(device)  # 定义VAE模型
    optimizer = torch.optim.Adam(vae.parameters(), lr=args.learning_rate)  # 使用Adam优化器
    logs = defaultdict(list)  # 记录日志
    for epoch in range(args.epochs):
        tracker_epoch = defaultdict(lambda: defaultdict(dict))  # 记录每个epoch的日志
        for iteration, (x, y) in enumerate(data_loader):
            x, y = x.to(device), y.to(device)  # 将数据放到设备上
            if args.conditional:  # 如果使用条件VAE
                recon_x, mean, log_var, z = vae(x, y)  # 前向传播 得到重构图像、均值、方差、隐变量
            else:
                recon_x, mean, log_var, z = vae(x)
            for i, yi in enumerate(y):
                id = len(tracker_epoch)  # 记录每个样本的id
                tracker_epoch[id]['x'] = z[i, 0].item()  # 记录每个样本的隐变量
                tracker_epoch[id]['y'] = z[i, 1].item()
                tracker_epoch[id]['label'] = yi.item()  # 记录每个样本的标签

            loss = loss_fn(recon_x, x, mean, log_var)  # 计算损失

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            logs['loss'].append(loss.item())

            if iteration % args.print_every == 0 or iteration == len(data_loader) - 1:  # 打印日志
                print("Epoch {:02d}/{:02d} Batch {:04d}/{:d}, Loss {:9.4f}".format(
                    epoch, args.epochs, iteration, len(data_loader) - 1, loss.item()))

                if args.conditional:  # 如果使用条件VAE
                    c = torch.arange(0, 10).long().unsqueeze(1).to(device)  # 定义标签
                    z = torch.randn([c.size(0), args.latent_size]).to(device)  # 定义隐变量
                    x = vae.inference(z, c=c)  # 推理得到重构图像
                else:
                    z = torch.randn([10, args.latent_size]).to(device)  # 定义隐变量
                    x = vae.inference(z)  # 推理得到重构图像

                plt.figure()  # 绘制图像
                plt.figure(figsize=(5, 10))  # 设置图像大小
                for p in range(10):   # 绘制10张图像
                    plt.subplot(5, 2, p + 1)  # 绘制5行2列图像  每行2张
                    if args.conditional:  # 如果使用条件VAE
                        plt.text(
                            0, 0, "c={:d}".format(c[p].item()), color='black',
                            backgroundcolor='white', fontsize=8)  # 显示标签
                    plt.imshow(x[p].view(28, 28).cpu().data.numpy())  # 显示重构图像
                    plt.axis('off')  # 关闭坐标轴

                if not os.path.exists(os.path.join(args.fig_root, str(ts))):
                    if not (os.path.exists(os.path.join(args.fig_root))):
                        os.mkdir(os.path.join(args.fig_root))
                    os.mkdir(os.path.join(args.fig_root, str(ts)))

                plt.savefig(  # 保存图片到文件
                    os.path.join(args.fig_root, str(ts),
                                 "E{:d}I{:d}.png".format(epoch, iteration)),
                    dpi=300)
                plt.clf()  # 清除图像
                plt.close('all')  # 关闭图像

        df = pd.DataFrame.from_dict(tracker_epoch, orient='index')  # 转换为DataFrame格式
        g = sns.lmplot(   # 绘制散点图
            x='x', y='y', hue='label', data=df.groupby('label').head(100), # 取前100个样本绘制散点图
            fit_reg=False, legend=True)  # 去掉拟合直线
        g.savefig(os.path.join(   # 保存散点图到文件
            args.fig_root, str(ts), "E{:d}-Dist.png".format(epoch)),
            dpi=300)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()  # 创建解析器
    parser.add_argument("--seed", type=int, default=0)  # 设置随机种子
    parser.add_argument("--epochs", type=int, default=10)  # 设置训练轮数
    parser.add_argument("--batch_size", type=int, default=64)  # 设置批大小
    parser.add_argument("--learning_rate", type=float, default=0.001)  # 设置学习率
    parser.add_argument("--encoder_layer_sizes", type=list, default=[784, 256])  # 设置编码器层数
    parser.add_argument("--decoder_layer_sizes", type=list, default=[256, 784])  # 设置解码器层数
    parser.add_argument("--latent_size", type=int, default=2)  # 设置隐变量维度
    parser.add_argument("--print_every", type=int, default=100)  # 设置打印频率
    parser.add_argument("--fig_root", type=str, default='figs')  # 设置图片保存路径
    parser.add_argument("--conditional", action='store_true')  # 设置是否使用条件VAE

    args = parser.parse_args()  # 解析命令行参数保存到args中
    print(args)  # 打印参数
    main(args)
